package com.txz.framework.utils.auth;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * 自定义请求包装类，支持修改请求头
 */
public class ModifyHeaderRequestWrapper extends HttpServletRequestWrapper {

    // 存储修改后的请求头
    private final Map<String, String> customHeaders;

    public ModifyHeaderRequestWrapper(HttpServletRequest request) {
        super(request);
        this.customHeaders = new HashMap<>();
    }

    /**
     * 添加或修改请求头
     */
    public void setHeader(String name, String value) {
        customHeaders.put(name, value);
    }

    /**
     * 优先返回自定义的请求头，没有则返回原始请求头
     */
    @Override
    public String getHeader(String name) {
        String headerValue = customHeaders.get(name);
        if (headerValue != null) {
            return headerValue;
        }
        // 调用父类方法获取原始请求头
        return super.getHeader(name);
    }

    /**
     * 合并原始请求头和自定义请求头（避免丢失原始头信息）
     */
    @Override
    public Enumeration<String> getHeaderNames() {
        // 合并原始头和自定义头的名称
        Map<String, String> allHeaders = new HashMap<>();
        // 添加原始请求头
        Enumeration<String> originalHeaderNames = super.getHeaderNames();
        while (originalHeaderNames.hasMoreElements()) {
            String name = originalHeaderNames.nextElement();
            allHeaders.put(name, super.getHeader(name));
        }
        // 添加自定义请求头（覆盖同名原始头）
        allHeaders.putAll(customHeaders);
        return Collections.enumeration(allHeaders.keySet());
    }
}